{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7bb66df4",
   "metadata": {},
   "source": [
    "# Building Agents \n",
    " \n",
    "> Note: Optionally, see [these slides](https://docs.google.com/presentation/d/13c0L1CQWAL7fuCXakOqjkvoodfynPJI4Hw_4H76okVU/edit?usp=sharing) and [langgraph_101.ipynb](langgraph_101.ipynb) for context before diving into this notebook!\n",
    "\n",
    "We're going to build an email assistant from scratch, starting here with 1) the agent architecture (using [LangGraph](https://langchain-ai.github.io/langgraph/)) and following with 2) testing (using [LangSmith](https://docs.smith.langchain.com/)), 3) human-in-the-loop, and 4) memory. This diagram show how these pieces will fit together:\n",
    "\n",
    "![overview-img](img/overview.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19d34429",
   "metadata": {},
   "source": [
    "#### Load environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46c9f78e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv(\"../.env\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54a69e9a",
   "metadata": {},
   "source": [
    "## Tool Definition\n",
    "\n",
    "Let's start by defining some simple tools that an email assistant will use with the `@tool` decorator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f2b708ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "from datetime import datetime\n",
    "from pydantic import BaseModel\n",
    "from langchain_core.tools import tool\n",
    "\n",
    "@tool\n",
    "def write_email(to: str, subject: str, content: str) -> str:\n",
    "    \"\"\"Write and send an email.\"\"\"\n",
    "    # Placeholder response - in real app would send email\n",
    "    return f\"Email sent to {to} with subject '{subject}' and content: {content}\"\n",
    "\n",
    "@tool\n",
    "def schedule_meeting(\n",
    "    attendees: list[str], subject: str, duration_minutes: int, preferred_day: datetime, start_time: int\n",
    ") -> str:\n",
    "    \"\"\"Schedule a calendar meeting.\"\"\"\n",
    "    # Placeholder response - in real app would check calendar and schedule\n",
    "    date_str = preferred_day.strftime(\"%A, %B %d, %Y\")\n",
    "    return f\"Meeting '{subject}' scheduled on {date_str} at {start_time} for {duration_minutes} minutes with {len(attendees)} attendees\"\n",
    "\n",
    "@tool\n",
    "def check_calendar_availability(day: str) -> str:\n",
    "    \"\"\"Check calendar availability for a given day.\"\"\"\n",
    "    # Placeholder response - in real app would check actual calendar\n",
    "    return f\"Available times on {day}: 9:00 AM, 2:00 PM, 4:00 PM\"\n",
    "\n",
    "@tool\n",
    "class Done(BaseModel):\n",
    "      \"\"\"E-mail has been sent.\"\"\"\n",
    "      done: bool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2911c929-5c41-4dcd-9cc8-21a8ff82b769",
   "metadata": {},
   "source": [
    "## Building our email assistant\n",
    "\n",
    "We'll combine a [router and agent](https://langchain-ai.github.io/langgraph/tutorials/workflows/) to build our email assistant.\n",
    "\n",
    "![agent_workflow_img](img/email_workflow.png)\n",
    "\n",
    "### Router\n",
    "\n",
    "The routing step handles the triage decision. \n",
    "\n",
    "The triage router only focuses on the triage decision, while the agent focuses *only* on the response. \n",
    "\n",
    "#### State\n",
    "\n",
    "When building an agent, it's important to consider the information that you want to track over time. We'll use LangGraph's pre-built [`MessagesState` object](https://langchain-ai.github.io/langgraph/concepts/low_level/#messagesstate), which is a just dictionary with a `messages` key that appends messages returned by nodes [as its update logic](https://langchain-ai.github.io/langgraph/concepts/low_level/#reducers). However, LangGraph gives you flexibility to track other information. We'll define a custom `State` object that extends `MessagesState` and adds a `classification_decision` key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "692537ec-f09e-4086-81e4-9c517273b854",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import MessagesState\n",
    "\n",
    "class State(MessagesState):\n",
    "    # We can add a specific key to our state for the email input\n",
    "    email_input: dict\n",
    "    classification_decision: Literal[\"ignore\", \"respond\", \"notify\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6cd1647-6d58-4aae-b954-6a9c5790c20c",
   "metadata": {},
   "source": [
    "#### Triage node\n",
    "\n",
    "We define a python function with our triage routing logic.\n",
    "\n",
    "For this, we use [structured outputs](https://python.langchain.com/docs/concepts/structured_outputs/) with a Pydantic model, which is particularly useful for defining structured output schemas because it offers type hints and validation. The descriptions in the pydantic model are important because they get passed as part JSON schema to the LLM to inform the output coercion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8adf520b-adf5-4a7b-b7a8-b8c23720c03f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pydantic import BaseModel, Field\n",
    "from email_assistant.utils import parse_email, format_email_markdown\n",
    "from email_assistant.prompts import triage_system_prompt, triage_user_prompt, default_triage_instructions, default_background\n",
    "from langchain.chat_models import init_chat_model\n",
    "from langgraph.graph import END\n",
    "from langgraph.types import Command"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "2c2c2ff0-da93-4731-b5b6-0ccd59e0e783",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"\\n\\n< Role >\\nYour role is to triage incoming emails based upon instructs and background information below.\\n</ Role >\\n\\n< Background >\\n{background}. \\n</ Background >\\n\\n< Instructions >\\nCategorize each email into one of three categories:\\n1. IGNORE - Emails that are not worth responding to or tracking\\n2. NOTIFY - Important information that worth notification but doesn't require a response\\n3. RESPOND - Emails that need a direct response\\nClassify the below email into one of these categories.\\n</ Instructions >\\n\\n< Rules >\\n{triage_instructions}\\n</ Rules >\\n\""
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rich.markdown import Markdown\n",
    "Markdown(triage_system_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3a1ad2c-40a2-42d0-a4b8-7a25df825fad",
   "metadata": {},
   "outputs": [],
   "source": [
    "Markdown(triage_user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b0df31-b9d2-423f-ba07-67eb0643c2ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "Markdown(default_background)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b3ea767-6ac1-4562-8ca6-5fa451495786",
   "metadata": {},
   "outputs": [],
   "source": [
    "Markdown(default_triage_instructions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c54ae6a6-94d9-4160-8d45-18f4d29aa600",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RouterSchema(BaseModel):\n",
    "    \"\"\"Analyze the unread email and route it according to its content.\"\"\"\n",
    "\n",
    "    reasoning: str = Field(\n",
    "        description=\"Step-by-step reasoning behind the classification.\"\n",
    "    )\n",
    "    classification: Literal[\"ignore\", \"respond\", \"notify\"] = Field(\n",
    "        description=\"The classification of an email: 'ignore' for irrelevant emails, \"\n",
    "        \"'notify' for important information that doesn't need a response, \"\n",
    "        \"'respond' for emails that need a reply\",\n",
    "    )\n",
    "\n",
    "# Initialize the LLM for use with router / structured output\n",
    "llm = init_chat_model(\"openai:gpt-4.1\", temperature=0.0)\n",
    "llm_router = llm.with_structured_output(RouterSchema) \n",
    "\n",
    "def triage_router(state: State) -> Command[Literal[\"response_agent\", \"__end__\"]]:\n",
    "    \"\"\"Analyze email content to decide if we should respond, notify, or ignore.\"\"\"\n",
    "    \n",
    "    author, to, subject, email_thread = parse_email(state[\"email_input\"])\n",
    "    system_prompt = triage_system_prompt.format(\n",
    "        background=default_background,\n",
    "        triage_instructions=default_triage_instructions\n",
    "    )\n",
    "\n",
    "    user_prompt = triage_user_prompt.format(\n",
    "        author=author, to=to, subject=subject, email_thread=email_thread\n",
    "    )\n",
    "\n",
    "    result = llm_router.invoke(\n",
    "        [\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt},\n",
    "        ]\n",
    "    )\n",
    "    \n",
    "    if result.classification == \"respond\":\n",
    "        print(\"📧 Classification: RESPOND - This email requires a response\")\n",
    "        goto = \"response_agent\"\n",
    "        update = {\n",
    "            \"messages\": [\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"Respond to the email: \\n\\n{format_email_markdown(subject, author, to, email_thread)}\",\n",
    "                }\n",
    "            ],\n",
    "            \"classification_decision\": result.classification,\n",
    "        }\n",
    "        \n",
    "    elif result.classification == \"ignore\":\n",
    "        print(\"🚫 Classification: IGNORE - This email can be safely ignored\")\n",
    "        goto = END\n",
    "        update =  {\n",
    "            \"classification_decision\": result.classification,\n",
    "        }\n",
    "        \n",
    "    elif result.classification == \"notify\":\n",
    "        print(\"🔔 Classification: NOTIFY - This email contains important information\")\n",
    "        # For now, we go to END. But we will add to this later!\n",
    "        goto = END\n",
    "        update = {\n",
    "            \"classification_decision\": result.classification,\n",
    "        }\n",
    "        \n",
    "    else:\n",
    "        raise ValueError(f\"Invalid classification: {result.classification}\")\n",
    "    return Command(goto=goto, update=update)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "272d8715",
   "metadata": {},
   "source": [
    "We use [Command](https://langchain-ai.github.io/langgraph/how-tos/command/) objects in LangGraph to both update the state and select the next node to visit. This is a useful alternative to edges.\n",
    "\n",
    "### Agent\n",
    "\n",
    "Now, let's build the agent.\n",
    "\n",
    "#### LLM node\n",
    "\n",
    "Here, we define the LLM decision-making node. This node takes in the current state, calls the LLM, and updates `messages` with the LLM output. \n",
    "\n",
    "We [enforce tool use with OpenAI](https://python.langchain.com/docs/how_to/tool_choice/) by setting `tool_choice=\"required\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1e842b3c-06f5-440f-8159-995503ef3a99",
   "metadata": {},
   "outputs": [],
   "source": [
    "from email_assistant.tools.default.prompt_templates import AGENT_TOOLS_PROMPT\n",
    "from email_assistant.prompts import agent_system_prompt, default_response_preferences, default_cal_preferences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f69c6fc-70aa-48f1-8312-2b1818469a1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "Markdown(AGENT_TOOLS_PROMPT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9052fced-3fdb-4cd2-ac88-e2ccdce14e7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "Markdown(agent_system_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6f2c120f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect all tools\n",
    "tools = [write_email, schedule_meeting, check_calendar_availability, Done]\n",
    "tools_by_name = {tool.name: tool for tool in tools}\n",
    "\n",
    "# Initialize the LLM, enforcing tool use\n",
    "llm = init_chat_model(\"openai:gpt-4.1\", temperature=0.0)\n",
    "llm_with_tools = llm.bind_tools(tools, tool_choice=\"any\")\n",
    "\n",
    "def llm_call(state: State):\n",
    "    \"\"\"LLM decides whether to call a tool or not\"\"\"\n",
    "\n",
    "    return {\n",
    "        \"messages\": [\n",
    "            # Invoke the LLM\n",
    "            llm_with_tools.invoke(\n",
    "                # Add the system prompt\n",
    "                [   \n",
    "                    {\"role\": \"system\", \"content\": agent_system_prompt.format(\n",
    "                        tools_prompt=AGENT_TOOLS_PROMPT,\n",
    "                        background=default_background,\n",
    "                        response_preferences=default_response_preferences,\n",
    "                        cal_preferences=default_cal_preferences, \n",
    "                    )}\n",
    "                ]\n",
    "                # Add the current messages to the prompt\n",
    "                + state[\"messages\"]\n",
    "            )\n",
    "        ]\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f05d11a",
   "metadata": {},
   "source": [
    "#### Tool handler node\n",
    "\n",
    "After the LLM makes a decision, we need to execute the chosen tool. \n",
    "\n",
    "The `tool_handler` node executes the tool. We can see that nodes can update the graph state to capture any important state changes, such as the classification decision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "43eb6dc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tool_handler(state: State):\n",
    "    \"\"\"Performs the tool call.\"\"\"\n",
    "\n",
    "    # List for tool messages\n",
    "    result = []\n",
    "    \n",
    "    # Iterate through tool calls\n",
    "    for tool_call in state[\"messages\"][-1].tool_calls:\n",
    "        # Get the tool\n",
    "        tool = tools_by_name[tool_call[\"name\"]]\n",
    "        # Run it\n",
    "        observation = tool.invoke(tool_call[\"args\"])\n",
    "        # Create a tool message\n",
    "        result.append({\"role\": \"tool\", \"content\" : observation, \"tool_call_id\": tool_call[\"id\"]})\n",
    "    \n",
    "    # Add it to our messages\n",
    "    return {\"messages\": result}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4721dede",
   "metadata": {},
   "source": [
    "#### Conditional Routing\n",
    "\n",
    "Our agent needs to decide when to continue using tools and when to stop. This conditional routing function directs the agent to either continue or terminate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7c7cbea7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def should_continue(state: State) -> Literal[\"tool_handler\", \"__end__\"]:\n",
    "    \"\"\"Route to tool handler, or end if Done tool called.\"\"\"\n",
    "    \n",
    "    # Get the last message\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    \n",
    "    # Check if it's a Done tool call\n",
    "    if last_message.tool_calls:\n",
    "        for tool_call in last_message.tool_calls: \n",
    "            if tool_call[\"name\"] == \"Done\":\n",
    "                return END\n",
    "            else:\n",
    "                return \"tool_handler\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eb4ede8",
   "metadata": {},
   "source": [
    "#### Agent Graph\n",
    "\n",
    "Finally, we can assemble all components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f81df767",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph, START, END\n",
    "from email_assistant.utils import show_graph\n",
    "\n",
    "# Build workflow\n",
    "overall_workflow = StateGraph(State)\n",
    "\n",
    "# Add nodes\n",
    "overall_workflow.add_node(\"llm_call\", llm_call)\n",
    "overall_workflow.add_node(\"tool_handler\", tool_handler)\n",
    "\n",
    "# Add edges\n",
    "overall_workflow.add_edge(START, \"llm_call\")\n",
    "overall_workflow.add_conditional_edges(\n",
    "    \"llm_call\",\n",
    "    should_continue,\n",
    "    {\n",
    "        \"tool_handler\": \"tool_handler\",\n",
    "        END: END,\n",
    "    },\n",
    ")\n",
    "overall_workflow.add_edge(\"tool_handler\", \"llm_call\")\n",
    "\n",
    "# Compile the agent\n",
    "agent = overall_workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "617f6373-bf48-44c2-ba33-000c9f22b067",
   "metadata": {},
   "outputs": [],
   "source": [
    "# View\n",
    "show_graph(agent)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc8367c4",
   "metadata": {},
   "source": [
    "This creates a graph that:\n",
    "1. Starts with an LLM decision\n",
    "2. Conditionally routes to tool execution or termination\n",
    "3. After tool execution, returns to LLM for the next decision\n",
    "4. Repeats until completion or no tool is called\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2b3406d-496d-43c9-942e-c5ce7e3a8321",
   "metadata": {},
   "source": [
    "### Combine workflow with our agent\n",
    "\n",
    "We can combine the router and the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "697f2548-b5a5-4fb6-8aed-226369e53e25",
   "metadata": {},
   "outputs": [],
   "source": [
    "overall_workflow = (\n",
    "    StateGraph(State)\n",
    "    .add_node(triage_router)\n",
    "    .add_node(\"response_agent\", agent)\n",
    "    .add_edge(START, \"triage_router\")\n",
    ").compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dd6dcc4-6346-4d41-ae36-61f3fc83b7a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_graph(overall_workflow, xray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2091d5cc",
   "metadata": {},
   "source": [
    "This is a higher-level composition where:\n",
    "1. First, the triage router analyzes the email\n",
    "2. If needed, the response agent handles crafting a response\n",
    "3. The workflow ends when either the triage decides no response is needed or the response agent completes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "070f18a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "email_input = {\n",
    "    \"author\": \"System Admin <sysadmin@company.com>\",\n",
    "    \"to\": \"Development Team <dev@company.com>\",\n",
    "    \"subject\": \"Scheduled maintenance - database downtime\",\n",
    "    \"email_thread\": \"Hi team,\\n\\nThis is a reminder that we'll be performing scheduled maintenance on the production database tonight from 2AM to 4AM EST. During this time, all database services will be unavailable.\\n\\nPlease plan your work accordingly and ensure no critical deployments are scheduled during this window.\\n\\nThanks,\\nSystem Admin Team\"\n",
    "}\n",
    "\n",
    "# Run the agent\n",
    "response = overall_workflow.invoke({\"email_input\": email_input})\n",
    "for m in response[\"messages\"]:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a50ae0a-7bd1-4e69-90be-781b1e77b4dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "email_input = {\n",
    "  \"author\": \"Alice Smith <alice.smith@company.com>\",\n",
    "  \"to\": \"John Doe <john.doe@company.com>\",\n",
    "  \"subject\": \"Quick question about API documentation\",\n",
    "  \"email_thread\": \"Hi John,\\nI was reviewing the API documentation for the new authentication service and noticed a few endpoints seem to be missing from the specs. Could you help clarify if this was intentional or if we should update the docs?\\nSpecifically, I'm looking at:\\n- /auth/refresh\\n- /auth/validate\\nThanks!\\nAlice\"\n",
    "}\n",
    "\n",
    "# Run the agent\n",
    "response = overall_workflow.invoke({\"email_input\": email_input})\n",
    "for m in response[\"messages\"]:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f631f61f",
   "metadata": {},
   "source": [
    "## Testing with Local Deployment\n",
    "\n",
    "You can find the file for our agent in the `src/email_assistant` directory:\n",
    "\n",
    "* `src/email_assistant/email_assistant.py`\n",
    "\n",
    "You can test them locally in LangGraph Studio by running:\n",
    "\n",
    "```\n",
    "! langgraph dev\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12752016",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "Example e-mail you can test:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08ee005a",
   "metadata": {},
   "outputs": [],
   "source": [
    "{\n",
    "  \"author\": \"Alice Smith <alice.smith@company.com>\",\n",
    "  \"to\": \"John Doe <john.doe@company.com>\",\n",
    "  \"subject\": \"Quick question about API documentation\",\n",
    "  \"email_thread\": \"Hi John,\\nI was reviewing the API documentation for the new authentication service and noticed a few endpoints seem to be missing from the specs. Could you help clarify if this was intentional or if we should update the docs?\\nSpecifically, I'm looking at:\\n- /auth/refresh\\n- /auth/validate\\nThanks!\\nAlice\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d09e33b6",
   "metadata": {},
   "source": [
    "![studio-img](img/studio.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d195e21-f2c5-4762-a4f0-c8d7459df6d5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
